package com.ruoyi.rms.utils;

import weka.classifiers.trees.RandomForest;
import weka.core.Attribute;
import weka.core.DenseInstance;
import weka.core.Instance;
import weka.core.Instances;
import weka.core.converters.ConverterUtils;

import java.math.BigDecimal;
import java.util.ArrayList;

public class RandomForestUtils {

    public static Instances data;

    public ArrayList<Attribute> initAttribute() {
        // 创建数据集
        ArrayList<Attribute> attributes = new ArrayList<>();
        attributes.add(new Attribute("价格"));
        attributes.add(new Attribute("销量"));
        attributes.add(new Attribute("评分"));

        return attributes;
    }

    public Instance initInstance(ArrayList<Attribute> attributes) {

        // 创建空的Instances对象
        data = new Instances("data", attributes, 0);

        return new DenseInstance(3);
    }

    public void loadTrainData(ArrayList<Attribute> attributes, Instance instance, double value1, double value2, double value3) {

        // 添加数据到Instances对象中
        instance.setValue(attributes.get(0), value1);
        instance.setValue(attributes.get(1), value2);
        instance.setValue(attributes.get(2), value3);
        data.add(instance);
    }

    public Instances loadTestData(ArrayList<Attribute> attributes, double value1, double value2, double value3) {
        // 创建空的Instances对象
        Instances testData = new Instances("testData", attributes, 0);

        DenseInstance instance = new DenseInstance(3);
        instance.setValue(attributes.get(0), value1);
        instance.setValue(attributes.get(1), value2);
        instance.setValue(attributes.get(2), value3);
        testData.add(instance);

        return testData;
    }

    public RandomForest train(Instances data) {

        try {
            // 使用DataSource读取Instances对象
            ConverterUtils.DataSource source = new ConverterUtils.DataSource(data);
            Instances newData = source.getDataSet();
            if (newData.classIndex() == -1) {
                newData.setClassIndex(newData.numAttributes() - 1);
            }

            // 创建随机森林分类器对象
            RandomForest rf = new RandomForest();
            rf.setMaxDepth(10); // 设置最大深度

            // 训练随机森林分类器对象
            rf.buildClassifier(newData);
            return rf;
        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }

    public double test(RandomForest rf, Instances testData) {
        try {
            // 使用DataSource读取Instances对象
            ConverterUtils.DataSource source = new ConverterUtils.DataSource(testData);
            Instances newData = source.getDataSet();
            if (newData.classIndex() == -1) {
                newData.setClassIndex(newData.numAttributes() - 1);
            }

            // 使用训练好的随机森林分类器对象进行预测
            return rf.classifyInstance(newData.instance(0));
        } catch (Exception e) {
            e.printStackTrace();
        }

        return 0.0;
    }
}
